import os
import tables as tb
import fnmatch
import numpy as np
import math
import argparse
import matplotlib.pylab as plt
from records import recordtype 
import help_func
from xml.etree.ElementTree import ElementTree

def calcSampMean(data):
    pos = np.random.randint(data.size, size=data.size) 
    samps = data[pos]
    return np.mean(samps)

def bootStrap(data,dataSqr,size):
    meanData = np.mean(data)
    meanDataSqr = np.mean(dataSqr)
    fMean = meanDataSqr-size*(np.power(meanData,2))
    sampMeanData = np.array([calcSampMean(data) for i in range(100)])
    sampMeanDataSqr = np.array([calcSampMean(dataSqr) for i in range(100)])
    fData = sampMeanDataSqr-size*(np.power(sampMeanData,2))
    var = np.var(fData)
    mean, var = fMean, var
    return mean, var 

class OrganizeMaxEntData:      
    def readOmegaMat(self, max_mat, omega_mat_type, L=100):
        mxData = self.maxEntAlgData
        if (omega_mat_type == "VER0"):
            omegaMatP = np.array([float(x) / float(L) for x in range(L / 2 + 1)])
        else: 
            omegMatFileName = help_func.locateFirst("kPointListAmp.dat", self.maxEntAlgData.dir)         
            omegMatFile = open(omegMatFileName, 'r')
            omegaMatP = omegMatFile.read().strip().split(' ')
            if (omega_mat_type == "VER1"):                
                omegaMatP = np.array(map(float, omegaMatP))
                omegaMatP = np.sqrt(omegaMatP) / float(L)
            elif (omega_mat_type == "VER2"):
                omegaMatP = map(lambda x:x.split('/'), omegaMatP)
                omegaMatP = np.array([math.sqrt(float(x[0]) / float(x[1])) for x in omegaMatP])
            else:
                print "No OmegaMat Version"
                exit(1)
        if (not mxData.zeroFreq):
            omegaMatP = omegaMatP[1:]
        #omegaMatP = omegaMatP[np.where(omegaMatP < self.maxEntAlgData.omegaCutOff)]
        return omegaMatP
    def readMeasureFile(self, sim_file):
        last = len(self.maxEntAlgData.omegaMatP)
        h5_file = tb.openFile(sim_file, "r")
        mxData = self.maxEntAlgData
        error_level = self.maxEntAlgData.errorLevel  
        sf = 0 if mxData.zeroFreq else 1
        obs=h5_file.getNode("/", mxData.obs + "/Mean")[error_level, sf:(last + 1)]
        obsVar = h5_file.getNode("/", mxData.obs + "/Variance")[error_level, sf:(last + 1), sf:(last + 1)]
        numOfMeasure = h5_file.getNode("/", mxData.obs + "/NumOfMeasure")[error_level]
        if (numOfMeasure < 10):
            return None
        obsVar = obsVar / (float(numOfMeasure * (numOfMeasure - 1)))
        #	if (mxData.zeroFreq):
        #            mname="AMPLITUDE"
        #            msqrname="AMPLITUDE_SQR"
        #            obszero = h5_file.getNode("/", mname + "/MeanBins").read() 
        #            obszerosqr = h5_file.getNode("/", msqrname + "/MeanBins").read() 
        #            cBin = h5_file.getNode("/", mname + "/CurrentMeasureBin").read()
        #            obszero = obszero[0:cBin + 1]   
        #            obszerosqr = obszerosqr[0:cBin + 1]
        #	    h5_file.close()
        #        return {"OBS":obs, "OBS_VAR":obsVar,"ZEROFREQ":obszero,"ZEROFREQ_SQR":obszerosqr}
        #	else:
        h5_file.close()
        return {"OBS":obs, "OBS_VAR":obsVar}
    def averageData(self, allData):
        n = len(allData)
        if (n > 1):
            obsInvCov = map(lambda x:np.linalg.inv(x), (z["OBS_VAR"] for z in allData)) 
            obsAvCov = np.linalg.inv(reduce(lambda x, y:x + y, obsInvCov))
            obsAvMu = reduce(lambda x, y:x + y , map(lambda x, y: np.dot(x, y), obsInvCov, (z["OBS"] for z in allData)))
            obsAvMu = np.dot(obsAvCov, obsAvMu)
        else:
            obsAvMu = allData[0]["OBS"]  
            obsAvCov = allData[0]["OBS_VAR"]
        err = np.sqrt(np.abs(np.diag(obsAvCov)))
#	if (self.maxEntAlgData.zeroFreq):
#	    zerofreq=np.concatenate([z["ZEROFREQ"] for z in allData])
#	    zerofreqsqr=np.concatenate([z["ZEROFREQ_SQR"] for z in allData])
#	    size = self.maxEntAlgData.beta ** 3
#            meanzero,varzero=bootStrap(zerofreq,zerofreqsqr,size)
#	    frow=np.hstack([self.maxEntAlgData.zeroFreqVar,np.zeros(obsAvMu.size)])
#            orow=np.vstack([np.zeros((1,obsAvMu.size)),obsAvCov])
#            obsAvMu=np.hstack([self.maxEntAlgData.zeroFreqVal,obsAvMu])
#            obsAvCov=np.hstack([frow[:,None],orow])
#	    err=np.sqrt(np.abs(np.diag(obsAvCov)))	    	
        if (self.maxEntAlgData.momentMes):
            momentInvCov = map(lambda x:1. / x, (z["MOMENT_VAR"] for z in allData)) 
            momentAvInvCov = 1. / (reduce(lambda x, y:x + y, momentInvCov))
            momentAvMu = reduce(lambda x, y:x + y , map(lambda x, y: x * y, momentInvCov, (z["MOMENT"] for z in allData)))
            momentAvMu = momentAvMu * momentAvInvCov
            return {"OBS":obsAvMu, "OBS_VAR":obsAvCov, "MOMENT":momentAvMu, "MOMENT_VAR":momentAvInvCov}
        if not os.path.exists(self.maxEntAlgData.dump_dir):
    		os.makedirs(self.maxEntAlgData.dump_dir)
        outFile = self.maxEntAlgData.dump_dir +  "/av_data.h5"
        h5file = tb.openFile(outFile, "w")
        h5file.createArray("/", "chi_s", obsAvMu)
        h5file.createArray("/", "chi_s_cov", obsAvCov)
        h5file.close()  # Moment
        return {"OBS":obsAvMu, "OBS_VAR":obsAvCov}
    def readData(self, maxEntAlgData, omega_mat_type):
        self.maxEntAlgData = maxEntAlgData
        self.simData = help_func.SimData
        # Find Files
        simFiles = help_func.get_sim_files(maxEntAlgData.dir)
        # Read Omega Mat
        # Find G
        xmlFileName = help_func.locateFirst("sim.xml", self.maxEntAlgData.dir)
        tree=ElementTree()
        xmlFileData=tree.parse(xmlFileName)
        elemG=(xmlFileData.find("G")).text
        self.maxEntAlgData.G=elemG
        elemMu=(xmlFileData.find("MU")).text
        self.maxEntAlgData.MU=elemMu
        elemBeta=int((xmlFileData.find("L_LIST")).findall("ITEM")[2].text)
        print "G= ",elemG," MU ",elemMu, " BETA ",elemBeta
        self.maxEntAlgData.beta=elemBeta
        self.maxEntAlgData.omegaMatP = self.readOmegaMat(maxEntAlgData.dir, omega_mat_type, self.maxEntAlgData.beta)
        # Read Data
        allData = [self.readMeasureFile(x["hdf5File"]) for x in simFiles]
        allData = [x for x in allData if x != None]
        #L=simFiles[0]["OBS"].size
        #self.maxEntAlgData.omegaMatP=np.array([float(x) / float(L) for x in range(L / 2 + 1)])
        allData = self.averageData(allData)
        err = np.sqrt(np.abs(np.diag(allData["OBS_VAR"])))
        print self.maxEntAlgData.omegaMatP.shape, " " ,allData["OBS"].shape
        help_func.errobarXY(self.maxEntAlgData.omegaMatP,allData["OBS"], err, self.maxEntAlgData.dump_dir + "/av_amp.pdf")
        # Buld Kernal
        self.buildKernel(allData)
        return self.simData
    def buildKernel(self, allData):
        mxdata = self.maxEntAlgData        
        # OmegeP       
        self.maxEntAlgData.omegaP = np.linspace(0., mxdata.omegaMax, \
                                            mxdata.numOfOmega, endpoint=False) 
          
     #   self.maxEntAlgData.omegaP =mxdata.omegaMax-mxdata.omegaMax*(1-np.exp(-alpha*np.tan((math.pi/2.)* np.linspace(0.,1. , \
      #                                      mxdata.numOfOmega, endpoint=False))))             
        # Build Kernel
        [OmegaP, OmegaMatP] = np.meshgrid(self.maxEntAlgData.omegaP, self.maxEntAlgData.omegaMatP)
        # Chec keigs
        self.check_eigs()
        #kernel = -OmegaP * (np.sinh(OmegaP) / (np.cos(2 * math.pi * OmegaMatP) - np.cosh(OmegaP)))
        kernel = -(1./math.pi) *OmegaP * (np.sinh(OmegaP) / (np.cos(2 * math.pi * OmegaMatP) - np.cosh(OmegaP)))
        if not self.maxEntAlgData.zeroFreq:
            kernel[:,0]= -(1./math.pi)
        # Rotate
        # Assa's Projection.
        [Up, sp, Vp] = np.linalg.svd(kernel, full_matrices=False)
        projdata=np.abs(np.dot(Up.T,allData["OBS"]))
        svds=sp
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.hold(True)
        ax.plot(projdata)
        ax.plot(svds)
        ax.set_yscale('log')
        fig.savefig(mxdata.dump_dir + "/svd_proj.pdf")
        fig.clf()
        outFile = self.maxEntAlgData.dump_dir +  "/proj.h5"
        h5file = tb.openFile(outFile, "w")
        h5file.createArray("/", "projdata", projdata)
        h5file.createArray("/", "svds", svds)
        h5file.close()  # Moment
        #
        val, vec = np.linalg.eigh(allData["OBS_VAR"])
        print np.linalg.norm(np.dot(vec, np.dot(np.diagflat(val), vec.T)) - allData["OBS_VAR"], 2)        
        kp = np.dot(vec.T, kernel)
        datap = np.dot(vec.T, allData["OBS"]) 
        # Add First moment
        if (self.maxEntAlgData.momentMes):
            kmoment = self.maxEntAlgData.omegaP * (1 - np.exp(-mxdata.beta * self.maxEntAlgData.omegaP))
            kp = np.vstack((kp, kmoment))
            val = np.append(val, (2 ** 2) * allData["MOMENT_VAR"])
            datap = np.append(datap, 2 * allData["MOMENT"])
            print "Moment" , 2 * allData["MOMENT"]
        # Reweight   
        w = 1. / val
        #w /= np.sum(w)
        datap = np.multiply(datap , np.sqrt(w))
        kp = np.multiply(kp , np.sqrt(w[ :, np.newaxis]))    
        self.simData.kp=kp
        self.simData.datap=datap
                   
    def check_eigs(self):
        vals = []
        mxdata = self.maxEntAlgData
        simFiles = help_func.get_sim_files(mxdata.dir)
        e0=self.maxEntAlgData.errorLevel
        for e in range(4):
            self.maxEntAlgData = mxdata
            self.maxEntAlgData.errorLevel = e
            # Read Data
            allData = self.averageData(map(lambda x: self.readMeasureFile(x["hdf5File"]), simFiles))
            val, vec = np.linalg.eigh(allData["OBS_VAR"])
            vals.append(np.flipud(np.sort(val)))
        self.maxEntAlgData.errorLevel=e0
        allData = self.averageData(map(lambda x: self.readMeasureFile(x["hdf5File"]), simFiles))
        fig = plt.figure()
        ax = fig.add_subplot(111)
        for e, v in enumerate(vals):
            ax.semilogy(v, label=str(e))        
        handles, labels = ax.get_legend_handles_labels()
        import operator
        hl = sorted(zip(handles, labels),
        key=operator.itemgetter(1))
        handles2, labels2 = zip(*hl)
        ax.legend(handles2, labels2)
        fig.savefig(self.maxEntAlgData.dump_dir + "/eigs.pdf")
            
if __name__ == "__main__":
    # arguments
    parser = argparse.ArgumentParser(description='MaxEnt')
    parser.add_argument('-root_dir', action='store')
    parser.add_argument('-obs', action='store')
    parser.add_argument('-moment', action='store')
    parser.add_argument('-max_mat', type=float, action='store')
    parser.add_argument('-error_level', type=int, action='store')
    parser.add_argument('-multi_sim', action='store_true', default=False)
    args = parser.parse_args()
    sim_data = OrganizeMaxEntData(args.root_dir, args.obs, args.moment)
    sim_data.read_data(args.error_level, args.max_mat)
    sim_data.write_data(args.root_dir + "/MaxEntData.h5")
    

